from datasets import list_datasets, load_dataset, list_metrics, load_metric
from transformers import AutoTokenizer

brt_tkn = AutoTokenizer.from_pretrained("bert-base-cased")


# Print all the available datasets
print(len(list_datasets()))

dataset = load_dataset('imdb')
dataset['train'][100]
dataset['train'].description
dataset['train'].features

def generate_tokens_for_imdb(examples):
    return brt_tkn(examples["text"], padding="max_length", truncation=True)

tkn_datasets = dataset.map(generate_tokens_for_imdb, batched=True)

training_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(200))
evaluation_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(200))

from transformers import AutoModelForSequenceClassification

mdl = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)

from transformers import TrainingArguments

training_args = TrainingArguments(output_dir="imdb")

import numpy as np
from datasets import load_metric


mdl_metrics = load_metric("accuracy")


def calculate_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return mdl_metrics.compute(predictions=predictions, references=labels)

from transformers import TrainingArguments, Trainer

trng_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch", num_train_epochs=3)

Mdl_trainer = Trainer(
    model=model,
    args=trng_args,
    train_dataset=training_dataset,
    eval_dataset=evaluation_dataset,
    compute_metrics=calculate_metrics,
)

trainer.train()

trainer.save_model()

metrics = mdl_trainer.evaluate(evaluation_dataset)
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)



